import hydra
import wandb
import os

import numpy as np

from tqdm import tqdm
from io import BytesIO
from PIL import Image

from omegaconf import DictConfig

from ecgcmr.signal.sig_augmentations.FindEDES import FindEDES
from ecgcmr.signal.sig_augmentations.ECGClean import DenoiseECG

from ecgcmr.signal.sig_utils.plot_ecg import plot_long_ecg


def log_ecg(ecg, patient, rpeaks=None, tpeaks=None):
    ecg_image_buffer = BytesIO()
    plot_long_ecg(ecg.squeeze(), rpeaks=rpeaks, tpeaks=tpeaks,
                  leads_to_plot=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], display=False, save_to=ecg_image_buffer)
    ecg_image_buffer.seek(0)
    ecg_image = Image.open(ecg_image_buffer)
    wandb.log({f"Patient {patient} ECG": wandb.Image(ecg_image)})


@hydra.main(config_path="ecgcmr/conf", config_name="base")
def main(cfg: DictConfig):
    wandb.init(project='Logging')
    
    save_folder = 'saved_tensors/multimodal/ecg_ED_ES'

    for mode in ["train", "val", "test"]:
        if mode == 'train':
            patient_ids = np.load(cfg.dataset.paths.data_patient_ids_train, mmap_mode='r')
            data_ecg = np.load(cfg.dataset.paths.data_ecg_train, mmap_mode='c')
        elif mode == 'val':
            patient_ids = np.load(cfg.dataset.paths.data_patient_ids_val, mmap_mode='r')
            data_ecg = np.load(cfg.dataset.paths.data_ecg_val, mmap_mode='c')
        elif mode == 'test':
            patient_ids = np.load(cfg.dataset.paths.data_patient_ids_test, mmap_mode='r')
            data_ecg = np.load(cfg.dataset.paths.data_ecg_test, mmap_mode='c')

        denoising = DenoiseECG(
                sampling_rate=cfg.dataset.sampling_rate,
                method=cfg.dataset.augmentations.ecg.clean_ecg.method
            )
        
        find_ed_es = FindEDES(cfg=cfg.dataset.augmentations.ecg,
                              input_electrodes=cfg.dataset.input_electrodes,
                              sampling_rate=cfg.dataset.sampling_rate
                              )
        
        valid_patient_ids = []
        valid_rpeaks = []
        valid_tpeaks = []
        valid_ratios = []

        for i in tqdm(range(len(patient_ids))):
            patient = patient_ids[i]
            ecg = data_ecg[i]

            try:
                denoised_ecg = denoising(sample=ecg)
                info_dict = find_ed_es(sample=denoised_ecg)

                rpeaks = info_dict['rpeaks']
                tpeaks = info_dict['tpeaks']
                ratio = info_dict['ratio']

                try:
                    ED_end_duration = 100 / ratio
        
                    delta_step = int(ED_end_duration - 50)

                    possible_rpeaks = [rpeak for rpeak in rpeaks if rpeak + 2500 <= 5000]

                    if len(possible_rpeaks)>0:
                        valid_patient_ids.append(patient)
                        valid_rpeaks.append(rpeaks)
                        valid_tpeaks.append(tpeaks)
                        valid_ratios.append(ratio)
                    else:
                        print()
                        print(f'Patient {patient} not enough r peaks')
                        log_ecg(ecg=ecg, patient=patient, rpeaks=rpeaks, tpeaks=tpeaks)

                except Exception as e:
                    print()
                    print(f'Patient {patient} ratio is nan')
                    log_ecg(ecg=ecg, patient=patient, rpeaks=rpeaks, tpeaks=tpeaks)

            except Exception as e:
                print(f'Patient {patient}, {e}')
                log_ecg(ecg=ecg, patient=patient)

        np.save(f"{save_folder}/{mode}_patient_ids.npy", np.array(valid_patient_ids))
        np.save(f"{save_folder}/{mode}_rpeaks.npy", np.array(valid_rpeaks, dtype=object))
        np.save(f"{save_folder}/{mode}_tpeaks.npy", np.array(valid_tpeaks, dtype=object))
        np.save(f"{save_folder}/{mode}_ratios.npy", np.array(valid_ratios))


if __name__ == "__main__":
    main()
